#!/usr/bin/python3
# -*- coding:utf-8 -*-
# __author__ = '__Tony__'
from functools import wraps
from sqlalchemy.exc import SQLAlchemyError


def session_begin(func):
    @wraps(func)
    def wrapper(db, **kwargs):  # 指定一毛一样的参数
        #if not db.in_transaction():
        with db.begin():
            return func(db, **kwargs)

    return wrapper



def async_transaction(db_arg_name="db"):
    """
    异步事务装饰器
    :param db_arg_name: 在目标函数中接收数据库session的参数名称
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # 获取数据库session对象
            db = kwargs.get(db_arg_name)
            if not db:
                # 如果在kwargs中找不到，尝试从args中根据参数名查找
                import inspect
                sig = inspect.signature(func)
                params = list(sig.parameters.values())
                for i, param in enumerate(params):
                    if param.name == db_arg_name and i < len(args):
                        db = args[i]
                        break

            if not db:
                raise ValueError(f"无法找到数据库session参数: {db_arg_name}")

            # 开始事务
            async with db.begin():
                try:
                    result = await func(*args, **kwargs)
                    # 自动提交（SQLAlchemy会自动处理）
                    return result
                except SQLAlchemyError:
                    # 自动回滚
                    await db.rollback()
                    raise
                except Exception:
                    # 非数据库错误也需要回滚
                    await db.rollback()
                    raise
        return wrapper
    return decorator
